"""
Core type definitions for the Claude Code Theory framework.

This module defines the fundamental types and data structures used throughout
the framework, following the mathematical formalization from the theoretical work.
"""

using LinearAlgebra
using Distributions
using StaticArrays

# ============================================================================
# Core System Types
# ============================================================================

"""
    abstract type AbstractClaudeCodeComponent

Base type for all components in the Claude Code system.
"""
abstract type AbstractClaudeCodeComponent end

"""
    abstract type AbstractState

Base type for all state representations in the system.
"""
abstract type AbstractState end

"""
    abstract type AbstractAction

Base type for all actions that can be taken by the system.
"""
abstract type AbstractAction end

"""
    abstract type AbstractObservation

Base type for all observations that the system can receive.
"""
abstract type AbstractObservation end

# ============================================================================
# File System and Environment Types
# ============================================================================

"""
    FileSystemState

Represents the state of the file system including file contents,
directory structure, and permissions.
"""
struct FileSystemState
    files::Dict{String, String}  # filepath -> content
    directories::Set{String}     # directory paths
    permissions::Dict{String, UInt8}  # filepath -> permission bits
    modification_times::Dict{String, Float64}  # filepath -> timestamp
end

"""
    DependencyGraph

Represents the dependency relationships in the codebase.
"""
struct DependencyGraph
    nodes::Set{String}  # code entities (functions, classes, modules)
    edges::Set{Tuple{String, String}}  # (from, to) dependencies
    weights::Dict{Tuple{String, String}, Float64}  # dependency strengths
    node_attributes::Dict{String, Dict{String, Any}}  # metadata per node
end

"""
    ProcessState

Represents the state of running processes and environment variables.
"""
struct ProcessState
    running_processes::Dict{Int, String}  # pid -> process name
    environment_variables::Dict{String, String}
    system_resources::Dict{String, Float64}  # resource usage metrics
end

"""
    ToolConfiguration

Represents the configuration and availability of tools.
"""
struct ToolConfiguration
    available_tools::Set{String}
    tool_configurations::Dict{String, Dict{String, Any}}
    tool_versions::Dict{String, String}
    disabled_tools::Set{String}
end

"""
    ExternalEnvironment

Represents external environment factors.
"""
struct ExternalEnvironment
    network_connectivity::Bool
    system_load::Float64
    available_memory::Float64
    available_disk::Float64
    timestamp::Float64
end

"""
    EnvironmentalState <: AbstractState

Complete environmental state as defined in Definition 2.1 of the framework.
Represents s = (ℱ, 𝒟, 𝒫, 𝒯_s, ℰ).
"""
struct EnvironmentalState <: AbstractState
    filesystem::FileSystemState
    dependencies::DependencyGraph
    processes::ProcessState
    tools::ToolConfiguration
    external::ExternalEnvironment
    hash_value::UInt64  # for efficient state comparison
    
    function EnvironmentalState(fs, deps, procs, tools, ext)
        hash_val = hash((fs, deps, procs, tools, ext))
        new(fs, deps, procs, tools, ext, hash_val)
    end
end

# ============================================================================
# Action and Tool Types
# ============================================================================

"""
    SafetyLevel

Enumeration of safety levels for tool operations.
"""
@enum SafetyLevel begin
    SAFE = 1      # No risk of system modification
    LOW_RISK = 2  # Minimal risk, reversible changes
    MEDIUM_RISK = 3  # Moderate risk, requires validation
    HIGH_RISK = 4   # High risk, requires confirmation
    DANGEROUS = 5   # Should not be executed without explicit approval
end

"""
    ToolCategory

Categories of tools as defined in Definition 2.2.
"""
@enum ToolCategory begin
    FILE_TOOLS     # File system operations
    SEARCH_TOOLS   # Code search and analysis
    EDIT_TOOLS     # Code modification
    EXEC_TOOLS     # Code execution
    TEST_TOOLS     # Testing and validation
    BUILD_TOOLS    # Building and compilation
    GIT_TOOLS      # Version control
    WEB_TOOLS      # Web and network operations
end

"""
    Tool

Represents a specific tool with its properties and constraints.
"""
struct Tool
    name::String
    category::ToolCategory
    description::String
    parameter_schema::Dict{String, Type}  # parameter name -> type
    return_schema::Dict{String, Type}     # return field -> type
    safety_level::SafetyLevel
    execution_cost::Float64  # expected computational cost
    average_latency::Float64 # expected execution time
    concurrent_compatible::Bool  # can run concurrently with others
    dependencies::Set{String}  # required tools/resources
end

"""
    ToolAction <: AbstractAction

Represents a tool invocation with parameters and concurrency specification.
Following Definition 2.3: a = (t, θ, c).
"""
struct ToolAction <: AbstractAction
    tool::Tool
    parameters::Dict{String, Any}  # parameter values
    concurrency_spec::BitVector    # which other tools can run in parallel
    priority::Float64              # execution priority
    timeout::Float64               # maximum execution time
end

# ============================================================================
# Observation Types
# ============================================================================

"""
    ToolExecutionResult

Result from executing a tool, including outputs and metadata.
"""
struct ToolExecutionResult
    success::Bool
    output::Any
    error::Union{Nothing, Exception}
    stdout::String
    stderr::String
    return_code::Int
    execution_time::Float64
    memory_usage::Float64
    disk_io::Float64
    network_io::Float64
end

"""
    EnvironmentalChange

Represents changes observed in the environment.
"""
struct EnvironmentalChange
    changed_files::Set{String}
    new_processes::Set{Int}
    terminated_processes::Set{Int}
    resource_changes::Dict{String, Float64}
    timestamp::Float64
end

"""
    UserFeedback

Represents feedback and steering signals from users.
"""
struct UserFeedback
    feedback_type::Symbol  # :approval, :rejection, :modification, :clarification
    content::String
    confidence::Float64    # user confidence in feedback
    timestamp::Float64
    context::Dict{String, Any}  # additional context
end

"""
    ExecutionMetadata

Metadata about the execution environment and performance.
"""
struct ExecutionMetadata
    execution_id::String
    start_time::Float64
    end_time::Float64
    resource_usage::Dict{String, Float64}
    error_conditions::Vector{String}
    warnings::Vector{String}
end

"""
    SystemObservation <: AbstractObservation

Complete observation as defined in Definition 2.4: o = (o_tool, o_env, o_user, o_meta).
"""
struct SystemObservation <: AbstractObservation
    tool_result::Union{Nothing, ToolExecutionResult}
    environment_change::Union{Nothing, EnvironmentalChange}
    user_feedback::Union{Nothing, UserFeedback}
    metadata::ExecutionMetadata
end

# ============================================================================
# Context and Belief Types
# ============================================================================

"""
    CodeKnowledge

Represents accumulated knowledge about the codebase.
"""
struct CodeKnowledge
    function_signatures::Dict{String, String}
    class_hierarchies::Dict{String, Vector{String}}
    import_relationships::Dict{String, Set{String}}
    test_coverage::Dict{String, Float64}
    complexity_metrics::Dict{String, Float64}
    semantic_embeddings::Dict{String, Vector{Float64}}
end

"""
    ConversationHistory

Represents the history of interactions with the user.
"""
struct ConversationHistory
    messages::Vector{Tuple{Symbol, String, Float64}}  # (role, content, timestamp)
    decisions::Vector{Tuple{String, Any, Float64}}    # (decision, params, timestamp)
    feedback_history::Vector{UserFeedback}
    performance_history::Vector{Float64}
end

"""
    ContextState

Complete context state as defined in Definition 2.6.
"""
struct ContextState
    code_knowledge::CodeKnowledge
    conversation_history::ConversationHistory
    file_system_summary::Dict{String, Any}
    execution_results::Vector{ToolExecutionResult}
    meta_information::Dict{String, Any}
    total_tokens::Int  # current token usage
    creation_time::Float64
end

"""
    BeliefState

Compressed belief state representation for computational tractability.
Following Assumption 2.2: φ(b_t) ∈ ℝ^d.
"""
struct BeliefState
    mean_estimate::Vector{Float64}      # d-dimensional mean
    covariance_estimate::Matrix{Float64} # d×d covariance matrix
    confidence::Float64                  # overall confidence [0,1]
    evidence_count::Int                 # number of observations
    last_update::Float64                # timestamp of last update
end

# ============================================================================
# Quality and Performance Types
# ============================================================================

"""
    QualityMetrics

Multi-dimensional quality assessment as per Definition 2.9.
"""
struct QualityMetrics
    correctness::Float64    # probability of being correct
    readability::Float64    # code readability score
    efficiency::Float64     # computational efficiency
    maintainability::Float64 # ease of maintenance
    
    function QualityMetrics(c::Real, r::Real, e::Real, m::Real)
        # Ensure all metrics are in [0,1]
        new(clamp(c, 0.0, 1.0), clamp(r, 0.0, 1.0), 
            clamp(e, 0.0, 1.0), clamp(m, 0.0, 1.0))
    end
end

"""
    RiskAssessment

Risk assessment across multiple dimensions as per Definition 2.10.
"""
struct RiskAssessment
    break_probability::Float64      # P(code break | s, a)
    security_risk::Float64         # P(security vulnerability | s, a)
    data_loss_risk::Float64        # P(data loss | s, a)
    performance_degradation::Float64 # P(performance impact | s, a)
    composite_risk::Float64        # overall risk score
    
    function RiskAssessment(break_p::Real, sec_p::Real, data_p::Real, perf_p::Real)
        # Calculate composite risk assuming independence
        composite = 1.0 - (1.0 - break_p) * (1.0 - sec_p) * (1.0 - data_p) * (1.0 - perf_p)
        new(clamp(break_p, 0.0, 1.0), clamp(sec_p, 0.0, 1.0),
            clamp(data_p, 0.0, 1.0), clamp(perf_p, 0.0, 1.0), 
            clamp(composite, 0.0, 1.0))
    end
end

"""
    ResourceMetrics

Resource usage and performance metrics.
"""
struct ResourceMetrics
    execution_time::Float64         # seconds
    memory_usage::Float64          # bytes
    cpu_usage::Float64             # percentage
    disk_io::Float64               # bytes
    network_io::Float64            # bytes
    api_calls::Int                 # number of API calls made
    cost_estimate::Float64         # estimated monetary cost
end

"""
    MultiObjectiveUtility

Multi-objective utility function as per Equation 2.1.
U = w_Q · Quality - w_L · Latency - w_C · Cost - w_R · Risk
"""
struct MultiObjectiveUtility
    weights::SVector{4, Float64}  # [w_Q, w_L, w_C, w_R]
    quality::Float64
    latency::Float64
    cost::Float64
    risk::Float64
    utility_value::Float64
    
    function MultiObjectiveUtility(w::AbstractVector, q::Real, l::Real, c::Real, r::Real)
        weights = SVector{4, Float64}(w)
        utility = weights[1] * q - weights[2] * l - weights[3] * c - weights[4] * r
        new(weights, Float64(q), Float64(l), Float64(c), Float64(r), utility)
    end
end

# ============================================================================
# Task and System Types
# ============================================================================

"""
    TaskType

Types of tasks the system can handle.
"""
@enum TaskType begin
    EXPLORATION     # Codebase understanding
    IMPLEMENTATION  # Code generation/modification
    TESTING        # Test creation and execution
    DEBUGGING      # Error detection and fixing
    REFACTORING    # Code restructuring
    DOCUMENTATION  # Documentation generation
    ANALYSIS       # Code analysis and metrics
end

"""
    CodeTask

Represents a development task to be executed by the system.
"""
struct CodeTask
    description::String
    task_type::TaskType
    programming_language::String
    target_files::Vector{String}
    constraints::Dict{String, Any}
    priority::Float64
    deadline::Union{Nothing, Float64}
    context_requirements::Dict{String, Any}
    success_criteria::Vector{String}
end

"""
    TaskResult

Result of executing a development task.
"""
struct TaskResult
    success::Bool
    outputs::Dict{String, Any}
    quality_metrics::QualityMetrics
    resource_metrics::ResourceMetrics
    risk_assessment::RiskAssessment
    execution_trace::Vector{ToolAction}
    duration::Float64
    error_log::Vector{String}
end

# ============================================================================
# Utility Functions
# ============================================================================

"""
    Base.hash(state::EnvironmentalState)

Efficient hash function for environmental states.
"""
Base.hash(state::EnvironmentalState) = state.hash_value

"""
    Base.:(==)(s1::EnvironmentalState, s2::EnvironmentalState)

Equality comparison for environmental states.
"""
Base.:(==)(s1::EnvironmentalState, s2::EnvironmentalState) = s1.hash_value == s2.hash_value

"""
    utility_value(u::MultiObjectiveUtility)

Extract the computed utility value.
"""
utility_value(u::MultiObjectiveUtility) = u.utility_value

"""
    is_safe(action::ToolAction, threshold::Float64 = 0.1)

Check if an action is considered safe based on risk assessment.
"""
function is_safe(action::ToolAction, threshold::Float64 = 0.1)
    return action.tool.safety_level <= LOW_RISK
end

"""
    total_risk(assessment::RiskAssessment)

Get the composite risk value from a risk assessment.
"""
total_risk(assessment::RiskAssessment) = assessment.composite_risk

"""
    total_quality(metrics::QualityMetrics, weights::Vector{Float64} = [0.25, 0.25, 0.25, 0.25])

Compute weighted total quality score.
"""
function total_quality(metrics::QualityMetrics, weights::Vector{Float64} = [0.25, 0.25, 0.25, 0.25])
    return weights[1] * metrics.correctness + 
           weights[2] * metrics.readability + 
           weights[3] * metrics.efficiency + 
           weights[4] * metrics.maintainability
end